[KV Cache] support kv cache int8 per channel quant #1662
Conversation
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @Eviannn, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request significantly enhances the LLM compression framework by implementing INT8 per-channel quantization for the Key and Value cache. The changes involve strategic tensor reshaping to align with the per-channel strategy and updates to the quantization observer to correctly compute and apply quantization parameters across the appropriate dimension, ultimately aiming to improve the efficiency and memory footprint of LLMs.
Highlights
- KV Cache Quantization: This pull request introduces support for INT8 per-channel quantization for Key and Value states within the KV cache, enhancing memory efficiency and performance for LLMs.
- Tensor Reshaping for Per-Channel Quantization: Key and Value tensors are now reshaped from
[batch_size, num_heads, seq_len, head_dim]to[batch_size, seq_len, num_heads * head_dim]before quantization, and then reshaped back after de-quantization. This prepares the tensors for per-channel quantization along the combined head dimension. - Dynamic Quantization Dimension Selection: The
BaseObserverlogic has been updated to dynamically select the correct dimension for per-channel quantization. Specifically, for 3-dimensional KV cache tensors, quantization parameters are now calculated alongdim=2(thenum_heads * head_dimdimension), while other scenarios continue to usedim=0. - Scale and Zero Point Handling: Adjustments were made to
squeeze()the scale and zero point tensors when storing them, accommodating the output format of per-channel quantization where these parameters might have singleton dimensions.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces support for per-channel INT8 quantization for the Key-Value (KV) cache. The changes involve reshaping the key/value tensors before quantization and then reshaping them back after dequantization. Additionally, the observer logic is updated to handle the 3D tensors characteristic of the KV cache.
My review has identified a couple of significant issues. There is a critical typo in the observer code that would prevent the new functionality from working correctly. Furthermore, the tensor reshaping logic within the KV cache is applied unconditionally, which would likely break existing quantization strategies. I've provided detailed comments and suggestions to address these points.
src/llmcompressor/observers/base.py
Outdated
| # 1. dim=2 scenario: in kv cache quant scenario which is | ||
| # [batch_size, seq_len - residual_length, num_heads * head_dim] | ||
| # 2. dim=0 scenario: assume observed is transposed, because its the output, hence use dim 0 | ||
| dim = 2 if observed.dim == 3 else 0 |
There was a problem hiding this comment.
There appears to be a typo in the condition here. observed.dim refers to the method object itself, not its return value. To get the number of dimensions, the method needs to be called: observed.dim().
Comparing the method object to an integer (observed.dim == 3) will always evaluate to False, causing dim to be incorrectly set to 0 for 3D tensors. This will prevent per-channel quantization from working correctly for use cases like the KV cache.
dim = 2 if observed.dim() == 3 else 0| # reshape for per channel scenario | ||
| num_heads = key_states.shape[1] | ||
| head_dim = key_states.shape[-1] | ||
| # from [batch_size, num_heads, seq_len - residual_length, head_dim] | ||
| # to [batch_size, seq_len - residual_length, num_heads * head_dim] | ||
| key_states = key_states.transpose(1, 2).flatten(2) | ||
| value_states = value_states.transpose(1, 2).flatten(2) |
There was a problem hiding this comment.
The reshape logic for key_states and value_states is specific to the per-channel quantization strategy, as noted in the comments. However, it's applied unconditionally, which will likely break other KV cache quantization strategies (e.g., per-tensor).
This reshape block should be wrapped in a conditional check, for example:
if self.quantization_args.strategy == QuantizationStrategy.CHANNEL:
This change will introduce a variable scope issue for num_heads and head_dim, which are defined in this block but also needed for the reverse reshape. You'll need to refactor the update method to handle this, for instance by initializing num_heads and head_dim to None at the beginning of the function.
You will also need to add the following import at the top of the file:
from compressed_tensors.quantization.quant_args import QuantizationStrategy
There was a problem hiding this comment.
pertensor and perchannel are same.
Signed-off-by: evian <eviantai@u.nus.edu>
Signed-off-by: evian <eviantai@u.nus.edu>
SUMMARY:
kv cache quant int8 per channel is supported using this pr.
Besieds, compressed-tensors needs to be updated as well: neuralmagic/compressed-tensors#398
TEST PLAN:
specify type=int8 and strategy=channel in recipe